package com.central.datax.plugin.reader.hivereader;

import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.util.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.datanucleus.store.types.wrappers.SqlDate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.security.PrivilegedExceptionAction;
import java.sql.*;
import java.util.*;

/**
 * @author Tindy
 * @date 2022/3/7
 * @describe
 */
public class HiveHelper {
    private static final Logger LOG = LoggerFactory.getLogger(HiveHelper.class);
    public static final String ORG_APACHE_HIVE_JDBC_HIVE_DRIVER = "org.apache.hive.jdbc.HiveDriver";
    private Configuration originalConfig;
    private Connection conn;
    public HiveHelper(Configuration originalConfig) {
        this.originalConfig=originalConfig;
    }

    public Connection getConn() throws SQLException {
        if(this.conn!=null && !this.conn.isClosed()){
            return this.conn;
        }
        try {
            Class.forName(ORG_APACHE_HIVE_JDBC_HIVE_DRIVER);
            Boolean haveKerberos = this.originalConfig.getBool(Key.HAVE_KERBEROS, false);
            if(haveKerberos) {
                synchronized (Constant.KRB5CONF_LOCK) {
                    UserGroupInformation.setLoginUser(null);
                    System.clearProperty("java.security.krb5.conf");
                    String krb5ConfPath=this.originalConfig.getString(Key.KERBEROS_CONF_FILE_PATH);
                    System.setProperty("java.security.krb5.conf", krb5ConfPath);
                    org.apache.hadoop.conf.Configuration conf = new org.apache.hadoop.conf.Configuration();
                    conf.set("hadoop.security.authentication", "Kerberos");
                    UserGroupInformation.setConfiguration(conf);
                    String keytabPath = this.originalConfig.getString(Key.KERBEROS_KEYTAB_FILE_PATH);
                    String kerberosPrincipal=this.originalConfig.getString(Key.KERBEROS_PRINCIPAL);
                    UserGroupInformation ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(kerberosPrincipal, keytabPath);
                    this.conn= ugi.doAs(new PrivilegedExceptionAction<Connection>() {
                        @Override
                        public Connection run() throws Exception {
                            return DriverManager.getConnection(originalConfig.getString(Key.JDBC_URL));
                        }
                    });
                }
            }else {
                this.conn = DriverManager.getConnection(this.originalConfig.getString(Key.JDBC_URL),this.originalConfig.getString(Key.USERNAME),this.originalConfig.getString(Key.PASSWORD));
            }
        } catch (Exception e) {
            LOG.error(e.getMessage(), e);
        }
        return this.conn;
    }

    public void closeConn(){
        try {
            if(this.conn!=null && !this.conn.isClosed()){
                conn.close();
            }
        }catch (SQLException e){
            throw DataXException.asDataXException(HiveReaderErrorCode.CONN_CLOSE_ERROR,e);
        }
    }

    public Collection<List<String>> getPartitionContitions(int adviceNumber) {
        if(adviceNumber<=1) return Collections.EMPTY_LIST;
        Map<Integer,List<String>> result = new HashMap<>();
        List<String> contitions=new ArrayList<>();
        String table = this.originalConfig.getString(Key.TABLE);
        Statement statement = null;
        try {
            statement = getConn().createStatement();
            ResultSet resultSet = statement.executeQuery("show partitions " +table);
            while (resultSet.next()){
                String partition = resultSet.getString("partition");
                String contition=partition2Contition(partition);
                contitions.add(contition);
            }
        } catch (SQLException e) {
            LOG.error(e.getMessage());
        }
        int size=contitions.size();
        for(int i=1;i<=size;i++){
            List<String> v = result.getOrDefault(i % adviceNumber, new ArrayList<>());
            v.add(contitions.get(i-1));
            result.put(i % adviceNumber,v);
        }
        return result.values();
    }

    private String partition2Contition(String partition){
        String[] kvsStr = partition.split("/");
        List<String> contitions=new ArrayList<>();
        for (String kvStr : kvsStr) {
            String[] kv = kvStr.split("=");
            String k=kv[0];
            String v= null;
            try {
                v = URLDecoder.decode(kv[1],"utf-8");
            } catch (UnsupportedEncodingException e) {
                LOG.error(e.getMessage());
            }
            contitions.add(k+"='"+v+"'");
        }
        String result = StringUtils.join(contitions, " and ");
        if(StringUtils.isNotEmpty(result)){
            result="( "+result+" )";
        }
        return result;
    }


    public ResultSet query(String querySql, int fetchSize) throws SQLException {

        Statement stmt = getConn().createStatement(ResultSet.TYPE_FORWARD_ONLY,
                ResultSet.CONCUR_READ_ONLY);
        stmt.setFetchSize(fetchSize);
        return stmt.executeQuery(querySql);
    }
}
